# 最短路算法水流动画
# voyagerhui
# 20220518
# 在给定图中，找出起点到每个顶点的最短路距离
# 本程序模拟一个物理过程，将图的边看作可以流水的管道，从顶点0处开始源源不断地注水，
# 水将以恒定速度流经各条边，遇到新的顶点后也将继续流动，直到所有管道都充满了水
# 水流先流到的顶点，将是距离起点更近的顶点。同时每个顶点距离起点的最短路距离为
# distance = 水流速度(v)✖时间(t)

"""
图的顶点
0
1
2
3
4
5
6
7
图的边，顶点1 顶点2 边长/权重
0 1 2
0 2 8
0 3 1
1 4 1
1 2 6
3 2 7
3 6 9
2 4 4
2 5 2
2 6 2
5 4 3
5 6 4
5 7 6
7 4 9
7 6 2
"""

import os
import matplotlib.pyplot as plt

V = 0.2  # 水流速度


# 图的顶点
class Vertex:
    def __init__(self, i, x, y):
        self.i = i  # 顶点编号
        self.neighbors = []  # 邻边上的邻点与边长[vertex, weight]
        self.flowed = False

        self.x = x
        self.y = y

    def add_neighbor(self, v, d):
        self.neighbors.append([v, d])

    def set_flowed(self):
        self.flowed = True

    def is_flowed(self):
        return self.flowed

    def get_i(self):
        return self.i

    def get_xy(self):
        return self.x, self.y

    def get_neighbors(self):
        return self.neighbors


# 一条边上的水流
class Flow:
    def __init__(self, si, ei, d, x1, y1, x2, y2):
        self.start_i = si
        self.end_i = ei
        self.d = d   # 边的长度
        self.d_flowed = 0  # 水流流经的长度

        self.start_x = x1
        self.start_y = y1
        self.end_x = x2
        self.end_y = y2
        self.curr_x = x1
        self.curr_y = y1

    def flow(self):
        self.curr_x = self.start_x + (self.end_x - self.start_x) * (self.d_flowed / self.d)
        self.curr_y = self.start_y + (self.end_y - self.start_y) * (self.d_flowed / self.d)
        self.d_flowed += V

    def curr_xy(self):
        return self.curr_x, self.curr_y

    def get_si_ei(self):
        return self.start_i, self.end_i

    def is_arrived(self):
        return self.d_flowed >= self.d


def create_graph():
    # create vertex
    v0 = Vertex(0, 0, 3)
    v1 = Vertex(1, 1, 0)
    v2 = Vertex(2, 3, 3)
    v3 = Vertex(3, 0.5, 6)
    v4 = Vertex(4, 4, 0)
    v5 = Vertex(5, 4, 3)
    v6 = Vertex(6, 5, 5)
    v7 = Vertex(7, 6, 3)
    vs = [v0, v1, v2, v3, v4, v5, v6, v7]
    # create edge
    es = [[0, 1, 2], [0, 2, 8], [0, 3, 1], [1, 4, 1], [1, 2, 6],
          [3, 2, 7], [3, 6, 9], [2, 4, 4], [2, 5, 2], [2, 6, 2],
          [5, 4, 3], [5, 6, 4], [5, 7, 6], [7, 4, 9], [7, 6, 2]]
    for e in es:
        vs[e[0]].add_neighbor(e[1], e[2])
        vs[e[1]].add_neighbor(e[0], e[2])
    return vs, es


# 从顶点i向邻边管道开始flow
def flow_from_vertex(flows, vs, i):
    vs[i].set_flowed()
    for item in vs[i].get_neighbors():
        si = i
        ei = item[0]
        d = item[1]
        [x1, y1] = vs[si].get_xy()
        [x2, y2] = vs[ei].get_xy()
        f = Flow(si, ei, d, x1, y1, x2, y2)
        flows.append(f)


# 水流算法
def flow_graph(vs, es):
    res = [[0, 0]]  # 存储顶点:距离

    flows = []  # 边上的水流集合
    flow_xy = []  # 存储水流流经的坐标
    flowed_edge = []  # 完全流过的边[[v1,v2], ...]
    flow_from_vertex(flows, vs, 0)  # 从顶点0开始注水
    t = 0  # 0时刻
    while True:
        all_flowed = False
        for f in flows:
            if f.is_arrived():
                si, ei = f.get_si_ei()
                flowed_edge.append([si, ei])
                if not vs[ei].is_flowed():  # 第一次到达
                    res.append([ei, int(t * V)])  # 添加改顶点到起点的距离
                    flow_from_vertex(flows, vs, ei)  # 从新顶点开始流动
            else:
                all_flowed = True
                f.flow()  # 更新水流的流动
                [x, y] = f.curr_xy()
                flow_xy.append([x, y])
        # 绘图
        plot_current(vs, es, flowed_edge, flow_xy, t, res)
        t += 1
        if not all_flowed:
            break


def plot_current(vs, es, flowed_edge, flow_xy, t, res):
    plt.figure()
    # plot vertex
    for vertex in vs:
        [x, y] = vertex.get_xy()
        i = vertex.get_i()
        plt.scatter(x, y, s=50, marker='o', color='red')
        plt.text(x + 0.2, y + 0.1, str(i), color='red')
    # plot edges
    for e in es:
        vertex1 = vs[e[0]]
        vertex2 = vs[e[1]]
        [x1, y1] = vertex1.get_xy()
        [x2, y2] = vertex2.get_xy()
        plt.plot([x1, x2], [y1, y2], color='black')
        plt.text((x1 + x2) / 2 + 0.1, (y1 + y2) / 2 + 0.1, str(e[2]), color='black')
    # plot flowed edges
    for e in flowed_edge:
        vertex1 = vs[e[0]]
        vertex2 = vs[e[1]]
        [x1, y1] = vertex1.get_xy()
        [x2, y2] = vertex2.get_xy()
        plt.plot([x1, x2], [y1, y2], color='blue')
    # plot water
    for item in flow_xy:
        plt.scatter(item[0], item[1], marker='s', color='blue')
    # save
    s = "t="+str(t)+" v:d="
    for item in res:
        s += (str(item[0])+":"+str(item[1])+"|")  # 顶点:最短路距离
    plt.title(s)
    plt.savefig('./pic/' + str(t) + '.png')
    plt.clf()


# 将./pic目录下的图片合成mp4动画
def generate_mp4():
    '''
    全部图片需要相同尺寸，从0.png开始编号
    :return:
    '''
    import cv2
    img = cv2.imread('./pic/0.png')  # 获取图像的尺寸
    video_writer = cv2.VideoWriter('./spa_flow.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 1, (img.shape[1], img.shape[0]))
    for i in range(len(os.listdir('./pic/'))):
        video_writer.write(cv2.imread('./pic/' + str(i) + '.png'))
    video_writer.release()


# 将./pic目录下的图片合成gif动画
def generate_gif():
    import imageio
    frames = []
    for i in range(len(os.listdir('./pic/'))):
        frames.append(imageio.imread('./pic/' + str(i) + '.png'))
    imageio.mimsave('./spa_flow.gif', frames, 'GIF', duration=1)


if __name__ == "__main__":
    vs, es = create_graph()
    flow_graph(vs, es)
    generate_mp4()
    generate_gif()
